
import torch
from torch.utils.data.sampler import Sampler
from collections import defaultdict
import random


def create_groups(groups, k):

    group_samples = defaultdict(list)
    for sample_idx, group_idx in enumerate(groups):
        group_samples[group_idx].append(sample_idx)

    keys_to_remove = []
    for key in group_samples:
        if len(group_samples[key]) < k:
            keys_to_remove.append(key)
            continue

    for key in keys_to_remove:
        group_samples.pop(key)

    return group_samples


class PKSampler(Sampler):

    def __init__(self, groups, p, k):
        self.p = p
        self.k = k
        self.groups = create_groups(groups, self.k)

        # Ensures there are enough classes to sample from
        assert len(self.groups) >= p

    def __iter__(self):
        # Shuffle samples within groups
        for key in self.groups:
            random.shuffle(self.groups[key])

        group_samples_remaining = {}
        for key in self.groups:
            group_samples_remaining[key] = len(self.groups[key])

        # while len(group_samples_remaining) >= self.p:
        #     group_ids = list(group_samples_remaining.keys())
        #     selected_group_idxs = torch.multinomial(torch.ones(len(group_ids)), self.p).tolist()
        #     for i in selected_group_idxs:
        #         group_id = group_ids[i]
        #         group = self.groups[group_id]
        #         for _ in range(self.k):
        #             sample_idx = len(group) - group_samples_remaining[group_id]
        #             yield group[sample_idx]
        #             group_samples_remaining[group_id] -= 1
        #
        #         if group_samples_remaining[group_id] < self.k:
        #             group_samples_remaining.pop(group_id)


        while len(group_samples_remaining) >= self.p:
            group_ids = list(group_samples_remaining.keys())
            selected_group_idxs = torch.multinomial(torch.ones(len(group_ids)), self.p).tolist()

            # 新增：收集当前批次的所有样本索引
            batch_indices = []

            for i in selected_group_idxs:
                group_id = group_ids[i]
                group = self.groups[group_id]
                for _ in range(self.k):
                    sample_idx = len(group) - group_samples_remaining[group_id]
                    # 原逻辑：直接yield → 改为先存入列表
                    batch_indices.append(group[sample_idx])
                    group_samples_remaining[group_id] -= 1

                if group_samples_remaining[group_id] < self.k:
                    group_samples_remaining.pop(group_id)

            # 新增：打乱当前批次内的样本顺序
            random.shuffle(batch_indices)

            # 逐个返回打乱后的索引
            for idx in batch_indices:
                yield idx
